package cn.bugstack.ai.handler;

import cn.bugstack.ai.model.*;
import cn.bugstack.ai.service.*;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

/**
 * AI WebSocket 处理器
 * 处理 WebSocket 连接和消息
 * 
 * @author bugstack
 */
@Slf4j
@Component
public class AIWebSocketHandler implements WebSocketHandler {
    
    private final AIService aiService;
    private final SessionManager sessionManager;
    private final ObjectMapper objectMapper = new ObjectMapper();
    
    public AIWebSocketHandler(AIService aiService, SessionManager sessionManager) {
        this.aiService = aiService;
        this.sessionManager = sessionManager;
        log.info("✅ AIWebSocketHandler initialized");
    }
    
    /**
     * 连接建立后调用
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String sessionId = session.getId();
        sessionManager.addSession(sessionId, session);
        log.info("✅ WebSocket connected - SessionId: {}, RemoteAddress: {}", 
                sessionId, session.getRemoteAddress());
        
        // 发送连接成功消息
        sendMessage(session, AIResponse.system(sessionId, "连接成功！您可以开始对话了。"));
    }
    
    /**
     * 接收到消息时调用
     */
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String payload = message.getPayload().toString();
        log.debug("📨 Received message - SessionId: {}, Payload: {}", session.getId(), payload);
        
        try {
            // 解析请求
            AIRequest request = objectMapper.readValue(payload, AIRequest.class);
            request.setSessionId(session.getId());
            
            log.info("📨 Processing request - SessionId: {}, Type: {}, Content: {}", 
                    session.getId(), request.getType(), 
                    request.getContent() != null && request.getContent().length() > 50 
                            ? request.getContent().substring(0, 50) + "..." 
                            : request.getContent());
            
            // 根据类型处理请求
            if ("trace".equals(request.getType())) {
                handleTraceRequest(session, request);
            } else {
                handleChatRequest(session, request);
            }
            
        } catch (Exception e) {
            log.error("❌ Error processing message - SessionId: {}, Error: {}", 
                    session.getId(), e.getMessage(), e);
            sendMessage(session, AIResponse.error(session.getId(), 
                    "处理消息时出错: " + e.getMessage()));
        }
    }
    
    /**
     * 处理普通聊天请求
     */
    private void handleChatRequest(WebSocketSession session, AIRequest request) {
        aiService.streamChat(request).subscribe(
                response -> sendMessage(session, response),
                error -> {
                    log.error("❌ Error in chat stream - SessionId: {}, Error: {}", 
                            session.getId(), error.getMessage());
                    sendMessage(session, AIResponse.error(request.getSessionId(), 
                            "处理请求时出错: " + error.getMessage()));
                },
                () -> log.debug("✅ Chat stream completed - SessionId: {}", session.getId())
        );
    }
    
    /**
     * 处理溯源请求
     */
    private void handleTraceRequest(WebSocketSession session, AIRequest request) {
        aiService.chatWithTracing(request).subscribe(
                response -> sendMessage(session, response),
                error -> {
                    log.error("❌ Error in trace stream - SessionId: {}, Error: {}", 
                            session.getId(), error.getMessage());
                    sendMessage(session, AIResponse.error(request.getSessionId(), 
                            "处理溯源请求时出错: " + error.getMessage()));
                },
                () -> log.debug("✅ Trace stream completed - SessionId: {}", session.getId())
        );
    }
    
    /**
     * 发送消息到客户端
     */
    private void sendMessage(WebSocketSession session, AIResponse response) {
        try {
            if (session.isOpen()) {
                String jsonResponse = objectMapper.writeValueAsString(response);
                session.sendMessage(new TextMessage(jsonResponse));
                log.trace("📤 Sent message - SessionId: {}, Type: {}", 
                        session.getId(), response.getType());
            } else {
                log.warn("⚠️ Session is closed, cannot send message - SessionId: {}", 
                        session.getId());
            }
        } catch (Exception e) {
            log.error("❌ Error sending message - SessionId: {}, Error: {}", 
                    session.getId(), e.getMessage());
        }
    }
    
    /**
     * 连接关闭后调用
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String sessionId = session.getId();
        sessionManager.removeSession(sessionId);
        log.info("❌ WebSocket disconnected - SessionId: {}, Status: {} ({})", 
                sessionId, status.getCode(), status.getReason());
    }
    
    /**
     * 传输错误时调用
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) {
        log.error("❌ Transport error - SessionId: {}, Error: {}", 
                session.getId(), exception.getMessage(), exception);
    }
    
    /**
     * 是否支持部分消息
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
